import pandas as pd
from datetime import datetime


def get_quarter(date_obj):
    """
    根据日期对象（pandas.Timestamp）返回季度。
    :param date_obj: 随访日期（pandas.Timestamp 类型）
    :return: 季度（1到4）
    """
    # 如果是字符串，先转换为 pandas Timestamp
    if isinstance(date_obj, str):
        date_obj = pd.to_datetime(date_obj)

    month = date_obj.month
    if 1 <= month <= 3:
        return 1
    elif 4 <= month <= 6:
        return 2
    elif 7 <= month <= 9:
        return 3
    else:
        return 4


def error6(df, styles_df, column1, column2, column3):
    """
    检查每个季度的随访建议是否一致，如果一致则标红。

    :param df: 原始数据 DataFrame
    :param styles_df: 样式 DataFrame，用于设置 Excel 格式
    :param column1: 身份证号
    :param column2: 随访日期
    :param column3: 随访建议
    """
    try:
        # 比较规则函数
        def highlight_if_duplicate_suggestions(row):
            idcard = row[column1]
            visit_date = row[column2]
            follow_up_advice = row[column3]

            # 获取当前行的季度
            current_quarter = get_quarter(visit_date)

            # 找出相同身份证号的所有记录
            same_id_df = df[df[column1] == idcard]

            # 获取当前身份证号的季度随访数据
            for _, other_row in same_id_df.iterrows():
                # 不比较同一行，获取另一行的季度
                if other_row[column2] != visit_date:
                    other_quarter = get_quarter(other_row[column2])

                    # 判断是否在相同季度之间
                    if current_quarter == other_quarter:
                        # 检查随访建议是否一致
                        if follow_up_advice == other_row[column3]:
                            # 标红：如果相同季度内随访建议一致
                            styles_df.loc[row.name, [column3]] = 'background-color: red'
                            styles_df.loc[other_row.name, [column3]] = 'background-color: red'

            return [''] * len(row)

        # 应用样式
        df.apply(highlight_if_duplicate_suggestions, axis=1)
        print("[随访建议]已完成处理。")

    except Exception as e:
        print(f"随访建议处理时发生错误：{e}")
